import tensorflow as tf
import t3f
from functools import reduce

NUM_CLASSES = 10

def tt_inference(images, hidden1_units, hidden2_units):
	modes_images = [4,7,7,4]
	modes_hidden1_units = [8,8,8,8]
	modes_hidden2_units = [6,6,6,6]

	if reduce(lambda x,y:x*y,modes_images) != images.shape[-1].value or \
		reduce(lambda x,y:x*y,modes_hidden1_units) != hidden1_units or \
		reduce(lambda x,y:x*y,modes_hidden2_units) != hidden2_units:
		return None

	with tf.variable_scope('hidden1'):
		initializer = t3f.glorot_initializer([modes_images, modes_hidden1_units], tt_rank = 32)
		tt_weights = t3f.get_variable('tt_weights', initializer = initializer)
		biases = tf.get_variable('biases', shape = [hidden1_units])
		hidden1 = tf.nn.relu(t3f.matmul(images, tt_weights) + biases)

	with tf.variable_scope('hidden2'):
		initializer = t3f.glorot_initializer([modes_hidden1_units, modes_hidden2_units], tt_rank = 48)
		tt_weights = t3f.get_variable('tt_weights', initializer = initializer)
		biases = tf.get_variable('biases', shape = [hidden2_units])
		hidden2 = tf.nn.relu(t3f.matmul(hidden1, tt_weights) + biases)

	with tf.variable_scope('softmax_linear'):
		weights = tf.get_variable('weights', shape = [hidden2_units, NUM_CLASSES])
		biases = tf.get_variable('biases', shape = [NUM_CLASSES])
		logits = tf.matmul(hidden2, weights) + biases

	return logits


# to be done...
def tt_inference_part(images, hidden2_units):
	modes_images = [4,7,7,4]
	modes_hidden2_units = [6,6,6,6]

	if reduce(lambda x,y:x*y,modes_images) != images.shape[-1].value or \
		reduce(lambda x,y:x*y,modes_hidden2_units) != hidden2_units:
		return None

	with tf.variable_scope('hidden2'):
		initializer = t3f.glorot_initializer([modes_images, modes_hidden2_units], tt_rank = 750)
		tt_weights = t3f.get_variable('tt_weights', initializer = initializer)
		biases = tf.get_variable('biases', shape = [hidden2_units])
		hidden2 = tf.nn.relu(t3f.matmul(images, tt_weights) + biases)

	with tf.variable_scope('softmax_linear'):
		weights = tf.get_variable('weights', shape = [hidden2_units, NUM_CLASSES])
		biases = tf.get_variable('biases', shape = [NUM_CLASSES])
		logits = tf.matmul(hidden2, weights) + biases

	return logits
